import os
import pandas as pd
import warnings

from utils import calculate_depreciation_rate_multiple_methods
warnings.filterwarnings('ignore', category=UserWarning, module='openpyxl')


class CapitalStockData:
    def __init__(self):
        self.ind_name = self.load_ind_names()
        self.total_k = self.load_data('固定资本形成总额.xlsx', keep_round=1)
        self.totak_k_without_house = pd.DataFrame()
        self.total_k_ind = self.load_data('年度固定资产投资完成额-19行业-2003-2017.xlsx', keep_round=1)
        self.total_k_ind_yoy = self.load_data('年度固定资产投资完成额-19行业-同比-2004-2024.xlsx', keep_round=1)
        self.build_k_ind = self.load_data('年度固定资产投资完成额-建筑安装工程-19行业-2003-2017.xlsx', keep_round=1)
        self.build_k_ind_yoy = self.load_data('年度固定资产投资完成额-建筑安装工程-19行业-同比-2018至今.xlsx', keep_round=1)
        self.equipment_k_ind = self.load_data('年度固定资产投资完成额-设备工器具购置-19行业-2003-2017.xlsx', keep_round=1)
        self.equipment_k_ind_yoy = self.load_data('年度固定资产投资完成额-设备工器具购置-同比-19行业-2018至今.xlsx', keep_round=1)
        self.other_k_ind = self.load_data('年度固定资产投资完成额-其他费用-19行业-2003-2017.xlsx', keep_round=1)
        self.other_k_ind_yoy = self.load_data('年度固定资产投资完成额-其他费用-同比-2018至今.xlsx', keep_round=1)
    
        self.house_k = pd.read_excel('data/House/住宅投资.xlsx', sheet_name='house', index_col=0)
        self.adjust_house_k()

        self.df_stock_2002 = pd.read_excel('data/kdata/2002年初始资本存量计算.xlsx', sheet_name='输出', index_col=0)
        self.df_stock_2002.index = [2002]
        self.price_index = pd.read_excel('data/kdata/固定资产投资价格指数1990至2019.xlsx', index_col=0)
        self.build_price_index = self.price_index.T.iloc[:,0]
        self.equipment_price_index = self.price_index.T.iloc[:,1]
        
        self.build_before_2018 = pd.DataFrame()
        self.equipment_before_2018 = pd.DataFrame()
        
        self.build_to_latest = pd.DataFrame()
        self.equipment_to_latest = pd.DataFrame()
        
        self.build_dic = None # 建筑安装工程 ，当年价格
        self.equipment_dic = None
        self.be_dic = None # 建筑安装工程 + 设备工器具购置
        
        self.build_dic_with_lasted_price = None # 建筑安装工程 ，2024年价格
        self.equipment_dic_with_lasted_price = None
        self.be_dic_with_lasted_price = None # 建筑安装工程 + 设备工器具购置 ，2024年价格
        
        self.stock_of_k = pd.DataFrame()
    
    def adjust_house_k(self):
        self.house_k.index = pd.to_datetime(self.house_k.index)
        # 把房地产业的建筑安装工程和设备工器具空间掉
        be_house = (self.build_k_ind + self.equipment_k_ind)['房地产']
        df = pd.concat([self.house_k, be_house], axis=1)
        df.dropna(inplace=True, how='any', axis=0)
        df.index =[s for s in df.index if s.year <=2024] # 防止更新到2025年数据后ratio会变
        # 计算放缩比例
        ratio = df.iloc[:,0].sum()/df.iloc[:,1].sum()
        # 放缩房地产业的建筑安装工程和设备工器具购置
        self.build_k_ind['房地产'] *= (1-ratio)
        self.equipment_k_ind['房地产'] *= (1-ratio)
        
        new_index = [s for s in self.total_k.index if s.year >= df.index[0].year]
        self.total_k = self.total_k.loc[new_index]
        
        idx = [s for s in self.total_k.index if s in self.house_k.index]
        a = self.total_k.loc[idx].iloc[:,0]
        b = self.house_k.loc[idx].iloc[:,0]
        self.totak_k_without_house = pd.DataFrame(a - b, columns=['剔除住宅后的固定资本形成总额'])
        print(self.house_k)

    def load_ind_names(self):
        ind_name = pd.read_excel('data/kdata/19行业名称.xlsx', index_col=0)
        dic1 = dict(zip(ind_name.iloc[:,0].tolist(), ind_name.iloc[:,-1].tolist()))
        dic2 = dict(zip(ind_name.iloc[:,1].tolist(), ind_name.iloc[:,-1].tolist()))
        return {**dic1, **dic2}
    
    def set_price_index(self, idx_se):
        year = 2024
        idx = idx_se.copy()
        idx.index = [s.year for s in idx.index]
        idx = idx / 100
        new_idx = {year: 1}
        for y in range(year-1, 1990, -1):
            new_idx[y] = new_idx[y+1] * idx[y]
        return pd.Series(new_idx)
        
    def load_data(self, path, keep_round=False):
        data = pd.read_excel(f'data/kdata/{path}', index_col=0)
        data = self.auto_transform_data(data)
        if keep_round:
            data = data.round(keep_round)
        data = self.change_col_name(data)
        return data

    def change_col_name(self, df):
        cols_map = {}
        for col in df.columns:
            cols_map[col] = col
            for c, v in self.ind_name.items():
                if c in col:
                    cols_map[col] = v
                    break 
        df.columns = [cols_map[col] for col in df.columns]
        return df
    
    @staticmethod
    def auto_transform_data(df):
        if isinstance(df.index[0],str):
            df = df.T
        new_idx = []
        for s in df.index:
            if pd.isna(s):
                continue
            elif 'Wind' in str(s):
                continue
            else:
                new_idx.append(s)
        # print(new_idx)
        df = df.loc[new_idx]
        df.index = pd.to_datetime(df.index)
        df.index.name = None
        df.columns.name = None
        return df

    @staticmethod
    def tab_data(a, b):
        a_tab = pd.DataFrame(a,index=b.index,columns=b.columns)
        a_and_b = pd.concat([a, a_tab])
        for year in range(2018, 2025):
            location = pd.to_datetime(f'{year}-12-31')
            last_year = pd.to_datetime(f'{year-1}-12-31')
            a_and_b.loc[location,:] = a_and_b.loc[last_year,:] * (1 + b.loc[location,:] / 100)
        return a_and_b.round(1)


    def tab_build_and_equipment_before2018(self):
        year = 2017 # 以2017年为数据基础
        correct_data = {}
        build_and_equipment_before2018 = self.build_k_ind + self.equipment_k_ind
        new_build = pd.DataFrame(
            index=self.build_k_ind.index,columns=self.build_k_ind.columns
            )
        new_equipment = pd.DataFrame(
            index=self.equipment_k_ind.index,columns=self.equipment_k_ind.columns
            )
        build_and_equipment_before2018 = self.build_k_ind + self.equipment_k_ind
        this_year = pd.to_datetime(f"{year}-12-31")
        bev = build_and_equipment_before2018.loc[this_year].sum()
        # real_v = self.total_k.loc[this_year].values[0] 
        # 这里剔除住宅后的固定资本形成总额
        real_v = self.totak_k_without_house.loc[this_year].values[0]
        coef = real_v / bev # 修正系数
        print(this_year, bev, real_v, coef)
        new_build.loc[this_year] = coef * self.build_k_ind.loc[this_year]
        new_equipment.loc[this_year] = coef * self.equipment_k_ind.loc[this_year]

        # 向前推，直到2003年
        for year in range(2016, 2002, -1):
            this_year = pd.to_datetime(f"{year}-12-31")
            after_year = this_year + pd.DateOffset(years=1)
            
            build_se_ratio = self.build_k_ind.loc[this_year] / self.build_k_ind.loc[after_year]
            equipment_se_ratio = self.equipment_k_ind.loc[this_year] / self.equipment_k_ind.loc[after_year]
            
            build_this_year = new_build.loc[after_year] * build_se_ratio
            equipment_this_year = new_equipment.loc[after_year] * equipment_se_ratio
            
            bev = (build_this_year + equipment_this_year).sum()
            
            # real_v = self.total_k.loc[this_year].values[0]
            # 这里剔除住宅后的固定资本形成总额
            real_v = self.totak_k_without_house.loc[this_year].values[0]

            coef = real_v / bev # 修正系数
            
            new_build.loc[this_year] = coef * build_this_year
            new_equipment.loc[this_year] = coef * equipment_this_year
            print(f"{this_year.year},求和为：{bev:.1f},固定资本形成为：{real_v:.1f},修正系数{coef:.4f}")
            correct_data[this_year] = coef
        
        self.build_before_2018 = new_build # 修正后的建筑安装工程
        self.equipment_before_2018 = new_equipment # 修正后的设备购置
        
        self.total_k.index = pd.to_datetime(self.total_k.index)
        new_bea_before_2018 = self.build_before_2018 + self.equipment_before_2018
        new_bea_before_2018.index = pd.to_datetime(new_bea_before_2018.index)
        sum_value = new_bea_before_2018.sum(axis=1)
        sum_value = sum_value.astype(float)

        # df_before_2018 = pd.concat([sum_value, self.total_k.iloc[:, 0]], axis=1)
        # 这里剔除住宅后的固定资本形成总额
        df_before_2018 = pd.concat([sum_value, self.totak_k_without_house.iloc[:, 0]], axis=1)
        df_before_2018 = df_before_2018.loc[sum_value.index]
        df_before_2018.columns = ['new_build_and_equipment_before2018', 'total_k']
        df_before_2018['check'] = df_before_2018['new_build_and_equipment_before2018'] - df_before_2018['total_k']
        if df_before_2018['check'].abs().max() > 1e-3:
            print(df_before_2018['check'].abs().max())
            raise ValueError("检查dataframe, total_k与加总后的建筑安装工程与设备工器具是否一致")
        print("核算一致")
        print(df_before_2018)
        return df_before_2018, correct_data

    def tab_build_and_equipment_after2018(self): 
        build_2024 = pd.concat(
                [   
                self.build_before_2018, 
                pd.DataFrame(
                    data=0, 
                    index=self.build_k_ind_yoy.index, 
                    columns=self.build_k_ind_yoy.columns
                    )
                ], axis=0)
        equipment_2024 = pd.concat(
                [   
                self.equipment_before_2018, 
                pd.DataFrame(
                    data=0, 
                    index=self.equipment_k_ind_yoy.index, 
                    columns=self.equipment_k_ind_yoy.columns
                    )
                ], axis=0)

        for year in range(2018, 2025):
            this_year = pd.to_datetime(f"{year}-12-31")
            last_year = pd.to_datetime(f"{year-1}-12-31")
            build_v = build_2024.loc[last_year, :] * (1 + self.build_k_ind_yoy.loc[this_year, :] / 100)
            build_v.astype(float)
            equipment_v = equipment_2024.loc[last_year, :] * (1 + self.equipment_k_ind_yoy.loc[this_year, :] / 100)
            equipment_v.astype(float)

            # ratio = ((build_v.sum() + equipment_v.sum()) / self.total_k.loc[this_year]).values[0]
            ratio = ((build_v.sum() + equipment_v.sum()) / self.totak_k_without_house.loc[this_year]).values[0]
            build_2024.loc[this_year, :] = build_v / ratio
            equipment_2024.loc[this_year, :] = equipment_v / ratio
            print(f"{year}年, 固定资产投资完成额: {build_v.sum():.2f}, 设备投资完成额: {equipment_v.sum():.2f}", ratio)
        
        self.build_to_latest = build_2024
        self.equipment_to_latest = equipment_2024
        
        be2024 = self.build_to_latest + self.equipment_to_latest
        be2024.index = pd.to_datetime(be2024.index)
        # df_2018_to_latest = pd.concat([ self.total_k, be2024.sum(axis=1)], axis=1)
        # 这里剔除住宅后的固定资本形成总额
        df_2018_to_latest = pd.concat([ self.totak_k_without_house, be2024.sum(axis=1)], axis=1)
        df_2018_to_latest = df_2018_to_latest.loc[be2024.index,:]
        df_2018_to_latest.columns = ['total_k','build_and_equipment']
        df_2018_to_latest['check'] = df_2018_to_latest['build_and_equipment'] - df_2018_to_latest['total_k']
        if df_2018_to_latest['check'].abs().max() > 1e-3:
            print(df_2018_to_latest['check'].abs().max())
            raise ValueError("检查dataframe, total_k与加总后的建筑安装工程与设备工器具是否一致")
        print("核算一致")
        print(df_2018_to_latest)
        self.build_to_latest.index = [s.year for s in self.build_to_latest.index]
        self.equipment_to_latest.index = [s.year for s in self.equipment_to_latest.index]
        return self.build_to_latest, self.equipment_to_latest

    def periodically_add_k(self, build_discount_year=55, equipment_discount_year=16):
        if self.build_to_latest is None or self.equipment_to_latest is None:
            df, correct_data = self.tab_build_and_equipment_before2018()
            build_to_latest, equipment_to_latest = self.tab_build_and_equipment_after2018()
        build_sum = self.build_before_2018.sum()
        equipment_sum = self.equipment_before_2018.sum()
        build_ratio = build_sum / (build_sum + equipment_sum)
        equipment_ratio = equipment_sum / (build_sum + equipment_sum)
        df_ratio_all = pd.DataFrame({
            '建筑安装工程': build_ratio,
            '设备工器具': equipment_ratio
        }).T
        se_build_ratio = df_ratio_all.iloc[0,:]
        se_equipment_ratio = df_ratio_all.iloc[1,:]
        se_build_ratio = se_build_ratio.astype(float)
        se_equipment_ratio = se_equipment_ratio.astype(float)
        
        df_stock_2002 = self.df_stock_2002.astype(float)
        se_stock_2002 = df_stock_2002.iloc[0,:]
        
        # 创建建筑安装工程2002年初始资本存量
        build_df_2002 = pd.DataFrame(se_stock_2002 * se_build_ratio)
        # 创建设备工器具2002年初始资本存量
        equipment_df_2002 = pd.DataFrame(se_stock_2002 * se_equipment_ratio)
        
        build_df_2002.columns = [2002]
        equipment_df_2002.columns = [2002]
        # 计算2002年至2023年的建筑安装工程资本存量
        depreciation_rate_build = calculate_depreciation_rate_multiple_methods(
            build_discount_year, 0.05, method='geometric'
            )
        depreciation_rate_equipment = calculate_depreciation_rate_multiple_methods(
            equipment_discount_year, 0.05, method='geometric'
        )
        build_price_index = self.set_price_index(self.build_price_index)
        equipment_price_index = self.set_price_index(self.equipment_price_index)
        
        build_dic = {2002: build_df_2002} 
        # key为年份，value为当年的资本存量的全息展示，比如说2003年的资本存量既有2002年的，也有2003年的
        equipment_dic = {2002: equipment_df_2002}
        be_dic = {2002: build_df_2002 + equipment_df_2002}
        
        for year in range(2003, 2025):
            keeper_build = build_dic[year-1] * (1 - depreciation_rate_build) #上年资产折旧到本年
            new_year_df_build = keeper_build.copy(deep=True) 
            new_year_df_build[year] = self.build_to_latest.loc[year,:] # 本年增量
            build_dic[year] = new_year_df_build
    
            keeper_equipment = equipment_dic[year-1] * (1 - depreciation_rate_equipment) #上年资产折旧到本年
            new_year_df_equipment = keeper_equipment.copy(deep=True) 
            new_year_df_equipment[year] = self.equipment_to_latest.loc[year,:] # 本年增量
            equipment_dic[year] = new_year_df_equipment
            
            be_dic[year] = build_dic[year] + equipment_dic[year]
        
        self.build_dic = build_dic
        self.equipment_dic = equipment_dic
        self.be_dic = be_dic
    
    def reset_price_of_k(self):
        build_dic_with_lasted_price = {}
        equipment_dic_with_lasted_price = {}
        be_dic_with_lasted_price = {}
        
        build_index = self.set_price_index(self.build_price_index)
        equipment_index = self.set_price_index(self.equipment_price_index)
        
        for year in range(2002, 2025):
            build_df = self.build_dic[year].copy(deep=True)
            equipment_df = self.equipment_dic[year].copy(deep=True)
            for col in range(2002, build_df.columns[-1]+1):
                build_df[col] = build_df[col] * build_index[col]
                equipment_df[col] = equipment_df[col] * equipment_index[col]
            build_dic_with_lasted_price[year] = build_df
            equipment_dic_with_lasted_price[year] = equipment_df
            be_dic_with_lasted_price[year] = build_df + equipment_df
        
        self.build_dic_with_lasted_price = build_dic_with_lasted_price
        self.equipment_dic_with_lasted_price = equipment_dic_with_lasted_price
        self.be_dic_with_lasted_price = be_dic_with_lasted_price
    
    def get_k_df(self):
        k_df = pd.DataFrame()
        for year in range(2002, 2025):
            # 每年的资本存量要累加起来
            k_df[year] = self.be_dic_with_lasted_price[year].T.sum(axis=0)
        k_df = k_df.astype(float).round(1)
        # 这里输出的资本存量k_df每一年的值均为存量数据，而不是增量数据
        return k_df

    def save_dic(self, ):
        # 把数据保存到excel
        with pd.ExcelWriter('output/k/建筑安装工程资本存量.xlsx') as writer:
            for name, df in self.build_dic.items():
                df.to_excel(writer, sheet_name=str(name))
        with pd.ExcelWriter('output/k/设备购置资本存量.xlsx') as writer:
            for name, df in self.equipment_dic.items():
                df.to_excel(writer, sheet_name=str(name))
        with pd.ExcelWriter('output/k/建筑安装工程加设备购置资本存量.xlsx') as writer:
            for name, df in self.be_dic.items():
                df.to_excel(writer, sheet_name=str(name))
        with pd.ExcelWriter('output/k/2024价格调整后建筑安装工程资本存量.xlsx') as writer:
            for name, df in self.build_dic_with_lasted_price.items():
                df.to_excel(writer, sheet_name=str(name))
        with pd.ExcelWriter('output/k/2024价格调整后设备购置资本存量.xlsx') as writer:
            for name, df in self.equipment_dic_with_lasted_price.items():
                df.to_excel(writer, sheet_name=str(name))
        with pd.ExcelWriter('output/k/2024价格调整后建筑安装工程加设备购置资本存量.xlsx') as writer:
            for name, df in self.be_dic_with_lasted_price.items():
                df.to_excel(writer, sheet_name=str(name))

    def how_diff_build_discount_rate_affect_k(self, equipment_discount_year=16):
        list_df = pd.DataFrame()
        for build_discount_year in range(30, 60):
            print(build_discount_year)
            discount_rate_build = calculate_depreciation_rate_multiple_methods(
                build_discount_year, 0.05, method='geometric'
            )
            self.periodically_add_k(
                build_discount_year=build_discount_year, equipment_discount_year=equipment_discount_year
                )
            self.reset_price_of_k()
            df = self.get_k_df()
            se = df.T.sum(axis=1).round(1)
            se.name = f"建筑折旧年限{build_discount_year}_折旧率{discount_rate_build:.4f}资本存量"
            # cap.save_dic()
            list_df[se.name] = se
        return list_df
    
    def how_both_discount_rate_affect_k(self, build_discount_year_range=(30, 60), equipment_discount_year_range=(5, 20)):
        bdy_range = range(*build_discount_year_range)
        edy_range = range(*equipment_discount_year_range)
        bdy_rate = [calculate_depreciation_rate_multiple_methods(build_discount_year, 0.05, method='geometric') 
                    for build_discount_year in bdy_range]
        edy_rate = [calculate_depreciation_rate_multiple_methods(equipment_discount_year, 0.05, method='geometric') 
                    for equipment_discount_year in edy_range]
        list_df = pd.DataFrame(index=bdy_range, columns=edy_range)
        for build_discount_year in bdy_range:
            print(build_discount_year)
            for equipment_discount_year in edy_range:
                self.periodically_add_k(
                    build_discount_year=build_discount_year, equipment_discount_year=equipment_discount_year
                    )
                self.reset_price_of_k()
                df = self.get_k_df()
                list_df.loc[build_discount_year, equipment_discount_year] = df.iloc[:,-1].sum()
        rate_df = list_df.copy(deep=True)
        rate_df.index = bdy_rate
        rate_df.columns = edy_rate
        return list_df, rate_df

if __name__ == '__main__':

    # 可变参数
    # 1. 建筑折旧年限
    # 2. 设备折旧年限
    # 3. 建筑安装工程价格指数
    # 4. 设备购置价格指数
    # 5. 初始年份（2002年）建筑安装工程资本存量
    # 6. 初始年份（2002年）设备购置资本存量

    cap = CapitalStockData()
    # 先递推2018年之前的分行业固定资本形成
    dddf, correct_data = cap.tab_build_and_equipment_before2018()
    # 再递推2018-2024年分行业固定资本形成
    build_2024, equipment_2024 = cap.tab_build_and_equipment_after2018()
    # 然后设定建筑折旧年限55年，设备折旧年限16年，计算资本存量
    cap.periodically_add_k(build_discount_year=55, equipment_discount_year=16)
    # 根据建筑安装工程价格指数和设备购置价格指数，计算2024年价格调整后的资本存量
    cap.reset_price_of_k()
    # cap.save_dic() # 保存2024年价格调整后的资本存量
    df = cap.get_k_df()
    # df.to_excel(f"output/k/2024价格调整后资本存量序列55_16.xlsx")
    # 如果有必要，通过设定不同的折旧率（建筑折旧年限和设备折旧年限），计算资本存量
    # a, b = cap.how_both_discount_rate_affect_k()